Sync upstream v2.17 - #99
Closed
Caozhou1995 wants to merge 84 commits into
Closed
Conversation
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
…A#2274) * Fix imports in test for deprecated jax.experimental.pjit Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix: Pass NamedSharding instead of PartitionSpec to compare_ops() so that when the in and out sharding is used to create a jitted function, it has the mesh info Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Signed-off-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kshitij Janardan Lakhani <klakhani@login-eos01.eos.clusters.nvidia.com>
* Support wheel build for cuda 13 Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fixes Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fixes for cu13 runtime, format Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Add documentation Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Better error handling Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * fix Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * fix jax sdist Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Modify function names Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…tization (NVIDIA#2270) * [JAX] Support recipe flags for disabling SR, RHT, and 2D quantization Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Fix issue with SR state being erased due to pytree handling of NVFP4Quantizer Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Add test for SR state preservation across VJP boundaries Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Fix sharding of SR rng state Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * update tolerances slightly now that SR is enabled Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Use hashlib for deterministic hashes across runs for SR Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * rename uses_rht on scaled tensors to has_applied_rht Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * add assert Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Move decision of whether to use RHT into helper.py and add dedicated RHT tests Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * fix use_rht attr usage Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * fix pure-jax rht usage criteria Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Adjust tolerances after rebase Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> --------- Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
Include TE core headers in build Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
* Added sm_120f to the build Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Change the arch specific handling Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Support for CUDA<12.9 Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Moved through the rest of the files Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Common cases Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Remove pure 100 from the list Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * CMake changes, (not yet working) Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Do not pass the arch-specific thing from build_tools Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Moved some of the files to arch-specific compilation Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix and also changing the order of compilation to hopefully get the compilation time lower Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Fix for the files overwriting custom compile properties Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Actually make this whole thing work Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add space to the error message Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com> * Apply suggestions from code review Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com> * Fixes from review Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * Changing the naming to be more intuitive Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add missing cassert include for device-side asserts Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
* add max_score for fused/unfused F16 non-CP Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * calculate max per head instead of max over all heads Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fused attn max_score shape Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert FE to github Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update FE to 1.15.0-rc Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix merge Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reduce ew kernels; fix causal masks; add more tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor fix to tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove logic for flash-attn Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * WIP: add CP support for p2p/a2a/all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor improvements of implementation/tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * WIP: add thd support Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add thd to UnfusedDPA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * more fixes for lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update to FE 1.15 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove unneeded changes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable unfused for thd + pad_between_seqs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable thd for unfused until bug is fixed Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix all gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * rename max_score to max_logit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable fused attn + thd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…DIA#2288) * Fix CI failures due to deterministic attention Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * some more cleanup Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fix debug test Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…2315) * Fix: Skip determinism tests for bprop for all sm >=100 Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add username to TODO Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Assert in fused attn bwd pass for sm100+ Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Fix attention backend and tests for sm120 Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Disable MLA only for backward Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…determinism on Blackwell (NVIDIA#2316) * Bump the min version expected to supported FP8 cs det on Blackwell Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Disable fused attn for cudnn < 9.14 for FP8 CS. Disable fused attn for cudnn < 9.18 for FP8 deterministic CS Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…NVIDIA#2322) Ensure JAX reference impl uses an accurate backend Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
# Description Add the FlagOS multi-chip backend for TransformerEngine Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --------- Co-authored-by: zhaoyinglia <ylzhao@baai.ac.cn>
# Description Fix import bugs Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
# Description Please include a brief summary of the changes, relevant motivation and context. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
…lagos-ai#4) # TransformerEngine-FL Plugin System ## Overview This PR implements a comprehensive multi-backend plugin system for TransformerEngine-FL, enabling support for multiple hardware vendors (NVIDIA, AMD, Hygon, etc.) while maintaining full API compatibility with the original `transformer_engine_torch`. **Core Philosophy**: A plugin-based backend system that allows hardware vendors to easily implement their own operator optimizations while preserving complete compatibility with the original TransformerEngine API. ## Key Features ### Full API Compatibility - Drop-in replacement for `transformer_engine_torch` - Switch backends via environment variables - Zero changes required to existing user code ### Multi-Backend Support | Backend | Description | Implementation | |---------|-------------|----------------| | **FlagOS (default)** | Triton-based cross-platform implementation | `backends/flagos/` | | **CUDA (vendor)** | Wraps original TransformerEngine C++ extensions | `backends/vendor/cuda/` | | **Reference** | Pure PyTorch fallback implementation | `backends/reference/` | ### Three-Tier Backend Selection ``` ┌─────────────────────────────────────────────────────────┐ │ 1. TE_FL_PER_OP (Per-operator override) [Highest] │ │ Example: TE_FL_PER_OP="rmsnorm_fwd=vendor:cuda" │ ├─────────────────────────────────────────────────────────┤ │ 2. TE_FL_PREFER (Global preference) │ │ Values: flagos / vendor / reference │ ├─────────────────────────────────────────────────────────┤ │ 3. Backend Priority (Intrinsic) [Lowest] │ │ Each implementation has a priority value │ └─────────────────────────────────────────────────────────┘ ``` ## Architecture ### Directory Structure ``` transformer_engine/plugin/core/ ├── __init__.py # Public API exports ├── types.py # Core types: BackendImplKind, OpImpl ├── registry.py # OpRegistry: stores all implementations ├── manager.py # OpManager: selects and calls implementations ├── policy.py # SelectionPolicy: backend selection rules ├── discovery.py # Plugin auto-discovery (entry_points, env) ├── builtin_ops.py # Registers all built-in backends ├── ops.py # TEFLModule: transformer_engine_torch compatible API ├── logger_manager.py # Logging utilities ├── _module_setup.py # Module aliasing setup ├── _build_config.py # Build-time configuration │ └── backends/ ├── flagos/ # FlagOS backend (Triton-based) │ ├── flagos.py # FlagOSBackend class │ ├── register_ops.py # Operator registration │ └── impl/ # Operator implementations │ ├── rmsnorm.py │ ├── gemm.py │ └── ... │ ├── vendor/ # Vendor backends │ └── cuda/ # NVIDIA CUDA backend │ ├── cuda.py # CUDABackend class │ └── register_ops.py │ └── reference/ # Reference backend (PyTorch) ├── reference.py # ReferenceBackend class ├── register_ops.py └── impl/ # Pure PyTorch implementations ``` ### Core Components | File | Description | |------|-------------| | `types.py` | Defines `BackendImplKind` (DEFAULT/VENDOR/REFERENCE) and `OpImpl` dataclass | | `registry.py` | `OpRegistry` - Central storage for all operator implementations | | `manager.py` | `OpManager` - Handles implementation selection, fallback, and execution | | `policy.py` | `SelectionPolicy` - Configurable rules for backend selection | | `discovery.py` | Auto-discovers plugins via `entry_points` or `TE_FL_PLUGIN_MODULES` | | `ops.py` | `TEFLModule` - Provides `transformer_engine_torch` compatible interface | ## Installation ### Build with CUDA support ```bash pip install --no-build-isolation -e . ``` ### Build without CUDA (FlagOS only) ```bash TE_FL_SKIP_CUDA=1 pip install --no-build-isolation -e . ``` ## Environment Variables ### Backend Selection | Variable | Description | Values | Default | |----------|-------------|--------|---------| | `TE_FL_PREFER` | Preferred backend type | `flagos` / `vendor` / `reference` | `flagos` | | `TE_FL_PREFER_VENDOR` | Prefer vendor (legacy) | `1` / `0` | `0` | | `TE_FL_STRICT` | Strict mode (no fallback) | `1` / `0` | `0` | ### Vendor Filtering | Variable | Description | Example | |----------|-------------|---------| | `TE_FL_ALLOW_VENDORS` | Allowed vendors (whitelist) | `nvidia,amd` | | `TE_FL_DENY_VENDORS` | Denied vendors (blacklist) | `vendor_a` | ### Per-Operator Configuration | Variable | Description | Example | |----------|-------------|---------| | `TE_FL_PER_OP` | Per-operator backend ordering | `rmsnorm_fwd=vendor:cuda\|default` | ### Plugin Discovery | Variable | Description | Example | |----------|-------------|---------| | `TE_FL_PLUGIN_MODULES` | Plugin modules to load | `my_plugin,another_plugin` | ### Build Configuration | Variable | Description | Values | Default | |----------|-------------|--------|---------| | `TE_FL_SKIP_CUDA` | Skip CUDA backend | `1` / `0` | `0` | | `CUDA_HOME` | CUDA installation path | `/usr/local/cuda` | Auto-detected | ### Logging | Variable | Description | Values | Default | |----------|-------------|--------|---------| | `TEFL_LOG_LEVEL` | Log level | `DEBUG` / `INFO` / `WARNING` / `ERROR` | `INFO` | ## Usage Examples ### Basic Usage (No Code Changes Required) ```python # Existing code works as-is import transformer_engine.pytorch as te # or import transformer_engine_torch as te ``` ### Register Custom Backend (In-tree) ```python from transformer_engine.plugin.core import ( OpRegistry, OpManager, OpImpl, BackendImplKind ) # 1. Define implementation def my_rmsnorm(input, weight, eps=1e-5, **kwargs): variance = input.pow(2).mean(-1, keepdim=True) return input * torch.rsqrt(variance + eps) * weight, torch.rsqrt(variance + eps) # 2. Register registry = OpRegistry() registry.register_impl(OpImpl( op_name="rmsnorm_fwd", impl_id="vendor.mybackend", kind=BackendImplKind.VENDOR, vendor="mybackend", fn=my_rmsnorm, priority=200, )) # 3. Call manager = OpManager(registry) output, rsigma = manager.call("rmsnorm_fwd", input, weight) ``` ### Register Custom Backend (Out-of-tree Plugin) Create a plugin package with `register(registry)` function: ```python # my_vendor_plugin/__init__.py from transformer_engine.plugin.core import OpImpl, BackendImplKind def my_rmsnorm(input, weight, eps=1e-5, **kwargs): # Your implementation ... def register(registry): """Called automatically by TE-FL""" registry.register_impl(OpImpl( op_name="rmsnorm_fwd", impl_id="vendor.myvendor", kind=BackendImplKind.VENDOR, vendor="myvendor", fn=my_rmsnorm, priority=200, )) ``` Load via environment variable: ```bash export TE_FL_PLUGIN_MODULES=my_vendor_plugin python your_script.py ``` ## Runtime Logs When running, you'll see logs indicating which backend is used: ``` [TE-FL manager.py:133 INFO] Registered impl_ids: ['default.flagos', 'reference.torch', 'vendor.cuda'] [TE-FL manager.py:390 INFO] Op 'rmsnorm_fwd' using 'default.flagos' (kind=default, vendor=None) [TE-FL manager.py:395 INFO] Op 'rmsnorm_fwd' switched from 'default.flagos' to 'vendor.cuda' (kind=vendor, vendor=CUDA) ``` ## Examples See `transformer_engine/plugins/examples/` for complete working examples: - `example_intree.py` - In-tree backend registration - `example_outtree.py` - Out-of-tree plugin registration Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --------- Co-authored-by: panpy <panpy@sugon.com>
# Description - Add missing __init__.py files to transformer_engine/plugin/core/backends/flagos/attention/ directory tree to fix import errors when accessing these modules as Python packages - Add comprehensive test suite (test_policy.py) covering the TE-FL scheduling policy system including: SelectionPolicy creation and configuration Environment variable parsing (TE_FL_PREFER, TE_FL_STRICT, etc.) Policy context managers Vendor filtering (allow/deny) Thread safety validation Minor code style improvements Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
…n fallback (flagos-ai#14) ## Summary This PR contains two major improvements: 1. **Register `get_attention_backend` function for all backends** (CUDA, FlagOS, Reference) - Added `get_attention_backend` implementation to all backend types - Ensures consistent attention backend selection across different hardware platforms 2. **Fix FlashAttention fallback mechanism** - Removed redundant `_called_impls` dictionary, replaced with simpler `_last_impl_id` class variable - Removed unused `_log_lock` threading lock - Simplified implementation tracking and logging logic - Reduced code complexity and memory overhead while maintaining full functionality ## Changes - Updated `FlashAttentionBase` class in `ops.py` to remove redundant implementation tracking - Added `get_attention_backend` registration to CUDA, FlagOS, and Reference backends - Fixed fallback logic in attention backend selection ## Test Plan - [x] Code builds successfully - [x] Existing tests pass - [x] Manual testing with different backend configurations ## Related Issues Fixes issues with FlashAttention fallback and improves backend consistency.
# Description fix nv shared lib bug [CUDA] Import failed: No module named 'transformer_engine_torch_nv' Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
# Description This pr add hygon backend for calling basic ops on hygon dcu. ## Type of change - [x] New feature (non-breaking change which adds functionality) ## Changes Please list the changes introduced in this PR: - Add a new `hygon` folder in `vendor` contains `__init__.py`, `hygon.py`, `register_ops.py` - Register hygon ops in `builtin_ops.py` # Requirements In order to use hygon backend, the following, the following requirements need to be met - The python package `transformer_engine_fl_hygon` needs to be installed # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --------- Signed-off-by: wenjh <wenjh@sugon.com>
…gos-ai#18) Add a flag that permanently enables flag_gems with a single switch, eliminating the need to call flag_gems.use_gems for every single operator. This removes significant registration overhead and improves end-to-end throughput. - When the flag is set, every operator’s implementation is forced to use flag_os/vendor; the default PyTorch reference backend is unavailable. - When the flag is not set, operators can freely switch among flag_os, vendor, and torch backends.
Unify the usage of the gems context - only enter or exit the context when switching between the flagos backend and the torch backend (or vice versa). - avoids the overhead of repeated enter/exit calls across multiple OPs.
## Summary - Support combined qkv_layout formats like `sbhd_sbhd_sbhd` by extracting the first part for layout conversion - Distinguish between standard 4D tensor format (sbhd/bshd) and true packed format (thd). For 4D tensors, directly convert layout like flagos backend does, instead of incorrectly trying to unpack ## Problem When using torch SDPA backend with `batch_size > 1`, the following error occurs: ``` ValueError: Unexpected 4D tensor shape torch.Size([4096, 4, 16, 128]). Expected [total_tokens, 1, num_heads, head_dim] ``` The original code incorrectly tried to unpack 4D tensors when `cu_seqlens` was provided, but 4D tensors in `sbhd`/`bshd` format should be handled with simple layout conversion (like flagos backend does). ## Test plan - [x] Tested with batch_size=4, verified no ValueError - [x] Results match flagos backend output
- Remove the flag_gems.use_gems() context to avoid context-switching overhead - Call flag_gems.xxx directly wherever possible.
# Description Add the new vendor backend METAX ## Type of change - [ ] New feature (non-breaking change which adds functionality) ## Changes Please list the changes introduced in this PR: - Add metax ops register - Add metax backend implementation - Register metax ops in builtin_ops.py ## Requirements - The module transformer_engine_torch_metax is needed, to use this module, need to install package transformer_engine_metax # Checklist: - [x] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes
…lagos-ai#23) ## Summary - flagos: Add multi_tensor_adam_param_remainder implementation - reference: Add multi_tensor_adam_param_remainder implementation - reference: Add context parallel support for Flash Attention - manager: Add cache mechanism with _impl_cache and _impl_cache_meta for conditional op selection ## Changes ### flagos backend - Implemented multi_tensor_adam_param_remainder operation for handling parameter remainders in multi-tensor Adam optimizer ### reference backend - Implemented multi_tensor_adam_param_remainder operation - Added context parallel support for Flash Attention implementation ### Core manager - Added cache mechanism using _impl_cache and _impl_cache_meta - Improved op selection with conditional caching based on policy fingerprint and epoch --------- Signed-off-by: wenone766 <wenone766@wenobug.com> Co-authored-by: wenone766 <wenone766@wenobug.com>
- Fix enum mismatch, between ```transformer_engine/plugin/core/ops.py``` and ```transformer_engine/common/include/transformer_engine/xxx.h```
# Description add Vendor KUNLUNXIN
Fix kunlunxin register errors
…g in the reference backend (flagos-ai#89) ## Summary This PR adds Ascend NPU support to the TE-FL plugin system through `torch_npu` and `transformer_engine_npu`, and fixes backward-path issues in the reference GEMM implementation. ## Changes ### Ascend NPU backend - Add automatic NPU availability detection and vendor-priority registration. - Add support for: - FlashAttention with SBHD, BSHD, and THD layouts - RMSNorm forward and backward - Generic and grouped GEMM - Multi-tensor scale and L2-norm operations - Add THD ↔ BSHD conversion operators. - Keep NPU dependencies lazily imported. ### Reference GEMM fixes - Fix output shape restoration for transposed inputs. - Do not add forward bias in backward mode. - Compute fused bias gradients. - Apply dGeLU using the saved forward activation. - Preserve correct alpha scaling and 3D input behavior. ## Testing Added coverage for: - FlashAttention forward/backward accuracy and causal masking - RMSNorm forward/backward - Generic and grouped GEMM - Multi-tensor and FP8 scale operations - Reference GEMM backward behavior Verified on Ascend 910C: ```text 46 passed ``` ## Deps It depends on TransformerEngineNPU. The package natively generated by TransformerEngineNPU is named transformer_engine. Relevant packaging logic needs to be modified so that the generated package is named transformer_engine_npu.
Move the temporary XTE TE-FL patch behavior into TE-FL native backend implementations. Register KunLunXin layernorm and GEMM operators, route attention backend selection through transformer_engine_klx_torch, and add reference GLU/DGLU fallback implementations. # Description Please include a brief summary of the changes, relevant motivation and context. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Change A - Change B # Checklist: - [ ] I have read and followed the [contributing guidelines](https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst) - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes
Summary This PR adds Ascend NPU Unit CI support for TransformerEngine-FL through torch_npu and the FlagOS backend. It extends the existing CI workflow to execute real TE workloads and selected shared PyTorch tests on Ascend 910C. Changes Ascend NPU testing Add real NPU coverage for: Linear forward and backward LayerNorm, RMSNorm, and LayerNormLinear LayerNormMLP GEMM, softmax, and multi-tensor operations Unfused Dot Product Attention and MultiheadAttention TransformerLayer debug and ONNX export paths Reuse selected portable sanity and numerics tests from the existing PyTorch suites. Distributed testing Add support for: Two-process HCCL execution TE Linear gradient synchronization Context Parallel utility tests Initial non-FP8 distributed numerical validation CI and coverage Add Ascend-specific Unit test entry points. Add raw and aggregated coverage collection. Fail explicitly when torch_npu or flag_gems is unavailable. Keep unsupported CUDA-specific features explicitly excluded. Testing Verified on Ascend 910C: PyTorch Unit test execution passed. PyTorch Debug passed. PyTorch ONNX passed. PyTorch Distributed is under validation. The current PyTorch Unit job failure occurred during coverage artifact upload after the test execution had passed. Limitations CUDA Graphs, Flash/Fused Attention, FP8, MXFP8, NVFP4, block scaling, TensorRT integration, and Integration tests are not included in the current Ascend Unit scope. --------- Co-authored-by: 1395976031 <1395976031@qq.com> Co-authored-by: BrianPei <kaworu228@gmail.com> Co-authored-by: wkhylyh-debug <wkhylyh@gmail.com>
…lagos-ai#92) ## Summary This PR adds a Hygon BW1000 CI baseline using the TE-FL reference backend and reorganizes the plugin tests into a backend-oriented structure under `tests/plugin`. The Hygon workflow validates the reference path only. It does not add or claim a native Hygon vendor backend. ## Changes - Add Hygon CI configuration, environment setup, and workflow entry. - Add Hygon unit, distributed smoke, ONNX smoke, and MCore integration tests. - Refactor common workflows to use platform configuration and setup scripts without chip-specific branches. - Move plugin tests from `transformer_engine/plugin/tests` to: - `tests/plugin/plugin` - `tests/plugin/backend/reference` - `tests/plugin/backend/flagos` - `tests/plugin/backend/npu` - `tests/plugin/backend/hygon` - Remove legacy plugin test files that were not collected by pytest. - Convert the FlagOS fused RoPE tests to standard pytest tests. - Add documentation for adding and running tests locally. - Use the unified 8-GPU runner labels. ## Hygon Baseline - Hardware: Hygon BW1000 - Backend policy: `TE_FL_PREFER=reference` - GEMM implementation: `reference.torch` - Runner label: `hg-8g-cicd-te` - Coverage enabled but not required - Debug tests are explicitly skipped when `nvdlfw_inspect` is unavailable ## Testing Validated on Hygon BW1000: - PyTorch unit tests - Plugin manager and policy tests - Reference backend tests - Distributed smoke tests - ONNX smoke tests - Coverage aggregation - Megatron-LM-FL MCore integration test All configured Hygon CI jobs passed. --------- Co-authored-by: wkhylyh-debug <wkhylyh@gmail.com>
## Summary Add a dedicated MUSA CI workflow for TransformerEngine-FL. ## Changes - Added MUSA hardware configuration and workflow entry points. - Added MUSA environment setup and runtime verification. - Verified the availability of the `transformer_engine_musa_torch` API and `vendor.musa`. - Verified that the representative `generic_gemm` dispatch selects `vendor.musa`. - Executed supported native TE tests using a MUSA-specific launcher. - Added a dedicated launcher for MUSA Megatron-LM integration tests. - Increased the timeout for shared unit tests from 60 minutes to 180 minutes. ## Test Organization - Native TE test adapter: `tests/plugin/backend/musa/run_native_tests.sh` - MUSA MCore integration test entry point: `tests/integration/musa/run_mcore.sh` Unsupported MUSA test cases are filtered out within the platform-specific launcher to ensure MUSA compatibility. ## Verification - YAML configuration parsing passed. - Bash syntax checks passed. - Python setup script checks passed. - `git diff --check` passed. - The final branch has been synchronized with `origin/musa-dev`. The MUSA backend implementation already exists in the upstream source code; this change provides the corresponding dedicated CI setup and test entry points. --------- Co-authored-by: canghaiX <1395976031@qq.com> Co-authored-by: BrianPei <kaworu228@gmail.com> Co-authored-by: wkhylyh-debug <wkhylyh@gmail.com> Co-authored-by: canghaiX <59075364+canghaiX@users.noreply.github.com>
|
wendell seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Caozhou1995
force-pushed
the
sync-upstream-v2.17
branch
from
August 10, 2026 04:03
fe32cde to
d29093a
Compare
added 6 commits
August 10, 2026 04:14
Resolved all 229 conflicts across P0/P1/P2 priorities: P0 (plugin system): No conflicts - fork's plugin/ directory fully preserved P1 (4 files): Manual merge preserving fork features + upstream additions 1. setup.py: Fork plugin build system + upstream NCCL EP support 2. transformer_engine/__init__.py: Fork plugin patches + upstream NCCL EP version check 3. transformer_engine/common/__init__.py: Fork skip_cuda_build() + plugin import 4. transformer_engine/pytorch/__init__.py: Fork torch_nv + upstream tensor/context utilities P2 (224 files): Accepted upstream version for core library files Upgrade path: v2.14 -> v2.17
- convert_host_pointers_to_tensor - get_device_pointer_for_data_and_scales These bindings were missing in upstream v2.17's pybind.cpp despite the corresponding C++ implementations existing in utils.cpp.
Applied 18 patches across 6 files to replace hardcoded 'cuda' device strings with TE_DEVICE_TYPE constant for multi-backend support. Changes: - utils.py: 4 patches - distributed.py: 3 patches - quantization.py: 6 patches - quantized_tensor.py: 1 patch - cpu_offload.py: 1 patch - jit.py: 3 patches All files passed Python AST syntax validation.
Added missing import for InstallCommand from setuptools.command.install to fix NameError during build.
These P2 files were git-added during Stage 3 merge but still contained conflict markers. Fixed by checking out the upstream v2.17 version for all affected files: - 43 source files (transformer_engine/common, pytorch, jax) - 10 test files - 4 doc/example files - 1 empty artifact removed (quantization_nvfp4.py renamed upstream) 64 files changed, accepting upstream version for all.
These files were generated during the v2.17 sync workflow for debugging and documentation purposes. They should not be part of the final PR.
Caozhou1995
force-pushed
the
sync-upstream-v2.17
branch
from
August 10, 2026 04:17
d29093a to
f7b076f
Compare
- ops.py: Add 15 new API declarations, update 4 existing signatures - cuda backend: Add 15 new methods, update 4 signatures in cuda.py and register_ops.py - All vendor backends (enflame/iluvatar/metax/musa/hygon): Add 19 new methods + OpImpl entries New APIs: splits_to_offsets_multi, copy_data_ptrs_to_device, bulk_allocate, create_empty_quantized_tensor, group_dequantize, get_grouped_gemm_setup_workspace_size, multi_tensor_pad_last_dim, multi_tensor_swizzle_scales_for_gemm_, multi_tensor_transpose_to_bhsd, cusolvermp_ctx_create, cusolvermp_ctx_destroy, newton_schulz, nvfp4_quantize_with_amax, nvfp4_group_quantize_with_amax, swizzle_scales_and_pack_ptrs_for_discrete_weights Modified signatures: group_quantize, bgrad_group_quantize (+tensor_offsets), clamped_swiglu, clamped_dswiglu (+glu_linear_offset)
The Stage 3 merge overwrote transformer_engine/__init__.py with upstream's version, losing the fork-specific TE_DEVICE_TYPE='cuda' default and the te_device_type() helper function. Multiple files depend on these: - transformer_engine/pytorch/utils.py - transformer_engine/pytorch/cpu_offload.py - transformer_engine/pytorch/distributed.py - transformer_engine/pytorch/jit.py - transformer_engine/debug/features/utils/stats_buffer.py - transformer_engine/pytorch/ops/fused/*.py
…Attention - Replace hardcoded device='cuda' with device=te_device_type() (2 locations) - Replace is_cuda checks with device.type == te_device_type() (3 locations) - These changes were lost during v2.17 upstream merge - Restores multi-backend compatibility from commit 4f54860
- triton/permutation.py: 12 te_device_type() calls restored - utils.py: 5 te_device_type() calls restored (_empty_tensor, normalize_device, torch_get_autocast_gpu_dtype)
…or, gemm, optimizers)
…ntion - Fix dot_product_attention.py: use 'from transformer_engine import te_device_type' instead of debug_state import - Add missing import in layernorm_mlp.py - Add missing import in quantization.py All imports now follow main branch convention: 'from transformer_engine import te_device_type'
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
19 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: